import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier,plot_tree
from sklearn.metrics import accuracy_score,precision_score,recall_score,f1_score,classification_report

# 1. 获取数据
data =pd.read_csv('titanic/train.csv')
data.info()
# 2.数据处理
# 2.1 获取特征+目标值
x = data[['Sex','Age','Pclass']]
print(x.head())
y = data['Survived']
# 2.2 热编码
x = pd.get_dummies(x)
print(x.head(10))
# 2.3 缺失值
x['Age'].fillna(x['Age'].mean(),inplace=True)
print(x.head(10))
# 2.4 数据集划分
x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=0.2,random_state=22)

# 3.特征工程(不需进行预处理)
# 4.模型训练
tree = DecisionTreeClassifier(criterion='gini',max_depth=3)
tree.fit(x_train,y_train)
# 5.模型预测
y_predict =tree.predict(x_test)
print(y_predict)
# 6.模型评估
print(accuracy_score(y_true=y_test, y_pred=y_predict))
print(precision_score(y_true=y_test, y_pred=y_predict))
print(recall_score(y_true=y_test, y_pred=y_predict))
print(f1_score(y_true=y_test, y_pred=y_predict))
print(classification_report(y_true=y_test, y_pred=y_predict))

# 树的可视化directX
plt.figure(figsize=(30,20))
plot_tree(tree,filled=True,feature_names=['Age','Pclass', 'Sex_female','Sex_male'],class_names=['died','no'])
plt.show()
